import * as vscode from "vscode";
import { runPythonResource } from "./subprocess";
import { buildPrompt, PromptData, invoke } from "./model";
import { getConfig } from "./config";
import { relative } from "path";

interface Language {
    FunctionSymbols: vscode.SymbolKind[];
    isFunctionSymbol(symbol: vscode.DocumentSymbol): boolean;
    getAllFnSymbols(): Promise<vscode.DocumentSymbol[]>;
    getActiveFnSymbol(): Promise<vscode.DocumentSymbol>;
    getActiveFnSignature(): Promise<string>;
    getActiveFnDocstring(): Promise<string>;
    getHintFromDocstring(doc: string): string;
    buildTypeContext(context: vscode.ExtensionContext): Promise<string>;
    buildCodeRetrieval(context: vscode.ExtensionContext): Promise<string>;
}

const Rust: Language = {
    FunctionSymbols: [
        vscode.SymbolKind.Function,
        vscode.SymbolKind.Method,
        vscode.SymbolKind.Constructor,
    ],
    isFunctionSymbol: (symbol: vscode.DocumentSymbol) =>
        Rust.FunctionSymbols.includes(symbol.kind),
    getAllFnSymbols: async () => {
        const allSymbols = (await vscode.commands.executeCommand(
            "vscode.executeDocumentSymbolProvider",
            vscode.window.activeTextEditor?.document.uri
        )) as vscode.DocumentSymbol[];
        const fnSymbols = allSymbols.filter(Rust.isFunctionSymbol);
        const implSymbols = allSymbols
            .filter((symbol) => symbol.kind === vscode.SymbolKind.Object)
            .flatMap((v, _, __) => v.children.filter(Rust.isFunctionSymbol));
        return fnSymbols.concat(implSymbols);
    },
    getActiveFnSymbol: async () => {
        const fnSymbols = await Rust.getAllFnSymbols();
        if (!fnSymbols || fnSymbols.length === 0) {
            throw new Error("No functions in this file");
        }
        const selection = vscode.window.activeTextEditor?.selection;
        if (!selection) {
            throw new Error("No active editor");
        }
        const selectedSymbols = fnSymbols.filter((symbol) =>
            symbol.range.intersection(selection)
        );
        if (!selectedSymbols || selectedSymbols.length === 0) {
            throw new Error("No functions are focused or selected");
        }
        if (selectedSymbols.length > 1) {
            const userChoice = await vscode.window.showQuickPick(
                selectedSymbols.map((v, _, __) => v.name),
                { title: "Select a function", canPickMany: false }
            );
            if (!userChoice) {
                throw new vscode.CancellationError();
            }
            return selectedSymbols.filter((v, _, __) => v.name === userChoice)[0];
        }
        return selectedSymbols[0];
    },
    getActiveFnSignature: async () => {
        const symbol = await Rust.getActiveFnSymbol();
        const detail = symbol.detail;
        const idx = detail.search("fn") + 2;
        return detail.slice(0, idx) + " " + symbol.name + detail.slice(idx);
    },
    getActiveFnDocstring: async () => {
        const sym = await Rust.getActiveFnSymbol();
        const range = new vscode.Range(sym.range.start, sym.selectionRange.start);
        const text = vscode.window.activeTextEditor?.document.getText(range);
        if (!text) {
            throw new Error("No active editor");
        }
        let lines = text.split("\n");
        lines.pop();
        return lines.map((v, _, __) => v.trim()).join("\n");
    },
    getHintFromDocstring: (doc: string) => {
        let lines = doc
            .split("\n")
            .map((v, _, __) => {
                v = v.trim();
                if (v.startsWith("///") || v.startsWith("//!")) {
                    v = v.substring(3).trim();
                } else if (v.startsWith("//")) {
                    v = v.substring(2).trim();
                }
                return v;
            })
            .filter((v, _, __) => v.length !== 0);
        if (lines.length > 0) {
            return lines[0];
        } else {
            return "";
        }
    },
    buildTypeContext: async (context: vscode.ExtensionContext) => {
        const path = getWorkspacePath();
        const src = getActiveFilePath();
        const sig = await Rust.getActiveFnSignature();
        return runPythonResource(context, "rust_context.py", path, src, sig);
    },
    buildCodeRetrieval: async (context: vscode.ExtensionContext) => {
        const ws_path = getWorkspacePath();
        const fn_symbol = await Rust.getActiveFnSymbol();
        const fn_name = fn_symbol.name;
        const fn_hint = Rust.getHintFromDocstring(
            await Rust.getActiveFnDocstring()
        );
        const fn_sig = await Rust.getActiveFnSignature();
        const fn_rel_path = relative(ws_path, getActiveFilePath());
        const start_line = fn_symbol.range.start.line.toString();
        const end_line = fn_symbol.range.end.line.toString();
        const embedding_model_path = getConfig<string>("embedding_model_path");
        return runPythonResource(
            context,
            "rust_retrieval.py",
            ws_path,
            fn_name,
            fn_hint,
            fn_sig,
            fn_rel_path,
            start_line,
            end_line,
            embedding_model_path
        );
    },
};

const getWorkspacePath = () => {
    const path = vscode.workspace.workspaceFolders?.[0].uri.fsPath;
    if (!path) {
        throw new Error("No workspace opened");
    }
    return path;
};

const getActiveFilePath = () => {
    const path = vscode.window.activeTextEditor?.document.uri.fsPath;
    if (!path) {
        throw new Error("No document opened");
    }
    return path;
};

const exceptionHandler = async (callback: () => Promise<string>, channel: vscode.OutputChannel) => {
    const statusBar = vscode.window.setStatusBarMessage('$(loading~spin) CatCoder Running');
    try {
        channel.clear();
        const code = await callback();
        channel.appendLine(code);
        vscode.window.setStatusBarMessage('Suggestion in Output > CatCoder', 8000);
    } catch (e: any) {
        if (e instanceof vscode.CancellationError) { }
        else if (e instanceof Error) {
            vscode.window.showErrorMessage('Encountered error. Details in in Output > CatCoder.');
            channel.appendLine(e.message);
            channel.appendLine(e.stack ?? "");
            console.log(e);
        }
    } finally {
        statusBar.dispose();
    }
};

export const rustImpl = (context: vscode.ExtensionContext, channel: vscode.OutputChannel) => async () =>
    await exceptionHandler(async () =>
        invoke(buildPrompt(await Promise.all([
            Rust.buildTypeContext(context),
            Rust.buildCodeRetrieval(context),
            Rust.getActiveFnDocstring(),
            Rust.getActiveFnSignature(),
        ])
            .then(
                (v) => new PromptData(...v),
                (err) => { throw err; }
            )
            .catch((err) => { throw err; })
        ))
        , channel);
